import pygame
import random

# 初始化Pygame
pygame.init()

# 游戏窗口尺寸
WIDTH = 300
HEIGHT = 600
BLOCK_SIZE = 30

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GRID_COLOR = (50, 50, 50)  # 新增网格线颜色
COLORS = [
    (0, 255, 255),  # 青色
    (255, 255, 0),  # 黄色
    (0, 255, 0),  # 绿色
    (255, 0, 0),  # 红色
    (0, 0, 255),  # 蓝色
    (255, 165, 0),  # 橙色
    (128, 0, 128)  # 紫色
]

# 方块形状定义（I, O, T, L, J, S, Z）
SHAPES = [
    [[1, 1, 1, 1]],  # I
    [[1, 1], [1, 1]],  # O
    [[0, 1, 0], [1, 1, 1]],  # T
    [[1, 0, 0], [1, 1, 1]],  # L
    [[0, 0, 1], [1, 1, 1]],  # J
    [[0, 1, 1], [1, 1, 0]],  # S
    [[1, 1, 0], [0, 1, 1]]  # Z
]


class Tetris:
    def __init__(self):
        self.screen = pygame.display.set_mode((WIDTH, HEIGHT))
        pygame.display.set_caption("俄罗斯方块")
        self.clock = pygame.time.Clock()
        self.grid = [[0 for _ in range(10)] for _ in range(20)]
        self.score = 0
        self.game_over = False
        self.spawn_new_piece()

    def spawn_new_piece(self):
        self.piece = {
            "shape": random.choice(SHAPES),
            "color": random.choice(COLORS),
            "x": 3,
            "y": 0,
            "rotation": 0
        }

    def rotate_piece(self):
        new_rotation = (self.piece["rotation"] + 1) % len(self.piece["shape"])
        rotated_shape = list(zip(*self.piece["shape"][::-1]))

        # 计算旋转后的形状边界
        min_x = min(x for row in rotated_shape for x, cell in enumerate(row) if cell)
        max_x = max(x for row in rotated_shape for x, cell in enumerate(row) if cell)
        width = max_x - min_x + 1

        # 尝试调整位置防止越界
        original_x = self.piece["x"]
        self.piece["x"] = max(0, min(10 - width, self.piece["x"]))

        # 检查碰撞
        if all(0 <= self.piece["x"] + x < 10 and 0 <= self.piece["y"] + y < 20
               for y, row in enumerate(rotated_shape)
               for x, cell in enumerate(row) if cell):
            self.piece["shape"] = [list(row) for row in rotated_shape]
            self.piece["rotation"] = new_rotation
        else:
            self.piece["x"] = original_x  # 回退位置

    def check_collision(self, dx, dy):
        for y, row in enumerate(self.piece["shape"]):
            for x, cell in enumerate(row):
                if cell:
                    new_x = self.piece["x"] + x + dx
                    new_y = self.piece["y"] + y + dy
                    if new_x < 0 or new_x >= 10 or new_y >= 20:
                        return True
                    if new_y >= 0 and self.grid[new_y][new_x]:
                        return True
        return False

    def lock_piece(self):
        for y, row in enumerate(self.piece["shape"]):
            for x, cell in enumerate(row):
                if cell:
                    self.grid[self.piece["y"] + y][self.piece["x"] + x] = self.piece["color"]

        # 消除完整行
        lines_cleared = 0
        for y in range(len(self.grid)):
            if all(self.grid[y]):
                del self.grid[y]
                self.grid.insert(0, [0] * 10)
                lines_cleared += 1

        self.score += lines_cleared * 100
        self.spawn_new_piece()

    def draw_grid(self):
        # 绘制背景网格
        for x in range(10):
            pygame.draw.line(self.screen, GRID_COLOR,
                             (x * BLOCK_SIZE, 0),
                             (x * BLOCK_SIZE, HEIGHT),
                             1)
        for y in range(20):
            pygame.draw.line(self.screen, GRID_COLOR,
                             (0, y * BLOCK_SIZE),
                             (WIDTH, y * BLOCK_SIZE),
                             1)

        # 绘制已固定的方块
        for y in range(20):
            for x in range(10):
                if self.grid[y][x]:
                    pygame.draw.rect(self.screen, self.grid[y][x],
                                     (x * BLOCK_SIZE, y * BLOCK_SIZE,
                                      BLOCK_SIZE - 1, BLOCK_SIZE - 1))

    def draw_piece(self):
        for y, row in enumerate(self.piece["shape"]):
            for x, cell in enumerate(row):
                if cell:
                    pygame.draw.rect(self.screen, self.piece["color"],
                                     ((self.piece["x"] + x) * BLOCK_SIZE,
                                      (self.piece["y"] + y) * BLOCK_SIZE,
                                      BLOCK_SIZE - 1, BLOCK_SIZE - 1))

    def draw_score(self):
        font = pygame.font.Font(None, 36)
        text = font.render(f"Score: {self.score}", True, WHITE)
        self.screen.blit(text, (10, 10))

    def run(self):
        fall_time = 0
        fall_speed = 500  # 毫秒

        while not self.game_over:
            self.screen.fill(BLACK)
            dt = self.clock.tick(30)
            fall_time += dt

            # 处理输入
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    self.game_over = True
                if event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_LEFT and not self.check_collision(-1, 0):
                        self.piece["x"] -= 1
                    if event.key == pygame.K_RIGHT and not self.check_collision(1, 0):
                        self.piece["x"] += 1
                    if event.key == pygame.K_DOWN:
                        fall_time = fall_speed  # 立即下落
                    if event.key == pygame.K_UP:
                        self.rotate_piece()

            # 自动下落
            if fall_time > fall_speed:
                if not self.check_collision(0, 1):
                    self.piece["y"] += 1
                else:
                    self.lock_piece()
                    # 检查游戏结束（修正后的条件）
                    if any(cell != 0 for cell in self.grid[0]):
                        self.game_over = True
                fall_time = 0

            # 绘制元素
            self.draw_grid()
            self.draw_piece()
            self.draw_score()

            pygame.display.flip()

        # 游戏结束显示
        font = pygame.font.Font(None, 72)
        text = font.render("Game Over!", True, WHITE)
        self.screen.blit(text, (WIDTH // 2 - 150, HEIGHT // 2 - 50))
        pygame.display.flip()
        pygame.time.wait(3000)
        pygame.quit()


if __name__ == "__main__":
    game = Tetris()
    game.run()